# -*- coding: utf-8 -*-
"""
@author: YuHaiyang

"""
import os
from pathlib import Path

import torch
from torch.utils.data import DataLoader
from torch.nn import functional as F
from nets.alex.alex_net import AlexNet
from nets.alex.loader import DataSetLoader


def print_hi(name):
    # Use a breakpoint in the code line below to debug your script.
    print(f'Hi, {name}')  # Press ⌘F8 to toggle the breakpoint.


# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    device = torch.device('cuda') if torch.cuda.is_available() else "cpu"

    home_path = os.environ.get("HOME")
    src = Path(home_path, "workspace", "dataset", "dc1000")
    loader: DataLoader = DataSetLoader(src).gen(batch_size=10)

    state_dict = torch.load("out/model0.pth", map_location=device)

    net = AlexNet()
    net.load_state_dict(state_dict)
    net.eval()

    print("loader:", enumerate(loader))
    with torch.no_grad():
        correct = 0.0
        test_loss = 0.0
        for index, (data, label) in enumerate(loader):
            data, label = data.to(device), label.to(device)
            output = net(data)
            loss = F.cross_entropy(output, label)
            pred = output.argmax(dim=1)
            print("acc:", label, ",pre:", pred)
            correct += pred.eq(label.view_as(pred)).sum().item()

        test_loss /= len(loader.dataset)
        print(
            "Test_average_loss : {:.4f} , Accuracy : {:.3f}%\n".format(test_loss, 100 * correct / len(loader.dataset)))
        acc = 100 * correct / len(loader.dataset)

